# -*- coding: utf-8 -*-
# 本程序用于视频中车辆行人等多目标检测跟踪
# @Time    : 2022/6/30 10:29
# @Software: PyCharm


#------------------------该文件用于实现车辆检测算法-----------------------------#

import os
import time

import cv2
import numpy as np

from sort import Sort
from speed import Car

#检测模型类
class Detector:
    # 初始化检测网络对象
    def __init__(self, model_path=None, video_path=None):
        '''
        Func:初始化检测模型对象
        Args:输入参数
            model_path:模型存放路径
            video_path:视频检测路径
        Param:
            self.LABELS:数据集标签(所有检测类别)
            self.net:检测网络
            self.ln:输出层
            self.tracker:追踪器
            self.filter_confidence:用于筛除置信度过低的识别结果
            self.threshold_prob:   用于NMS去除重复的锚框
        '''
        self.filter_confidence = 0.5  # 用于筛除置信度过低的识别结果
        self.threshold_prob = 0.3  # 用于NMS去除重复的锚框

        if model_path is None:
            model_path = "../yolo-obj"  # 模型文件的目录

        if video_path is None:
            video_path = '../video/car_chase_01.mp4'

        # 载入模型参数文件及配置文件
        # weightsPath = os.path.sep.join([model_path, "yolov4-tiny.weights"])
        # configPath = os.path.sep.join([model_path, "yolov4-tiny.cfg"])
        weightsPath = "".join(model_path)
        configPath = model_path.replace("weights", "cfg")

        # 载入数据集标签
        # # labelsPath = os.path.sep.join(["../yolo-obj", "coco.names"])
        # labelsPath = "./yolo-obj/coco.names"
        # self.LABELS = open(labelsPath).read().strip().split("\n")
        self.net = None # 检测网络
        # 从配置和参数文件中载入模型
        try:
            print("[INFO] 正在载入模型...")
            self.net = cv2.dnn.readNetFromDarknet(configPath, weightsPath)
            # 获取输出层的名字
            ln = self.net.getLayerNames() # 200,227,254
            self.ln = [ln[i[0] - 1] for i in self.net.getUnconnectedOutLayers()] # 82,94,106
            self.tracker = Sort()  # 实例化追踪器对象
        except:
            print("读取模型失败，请检查文件路径并确保无中文文件夹！")

    # 图像检测
    def run(self, frame_in):
        '''

        Args:
            frame: 图像帧

        Returns:
            dets: 图像检测结果
            boxes:追踪到的标记框
            indexIDs:追踪到的序号
            cls_IDs:踪到的类别

        '''
        # frame_in = frame.copy()
        # 将一帧画面读入网络
        blob = cv2.dnn.blobFromImage(frame_in, # 图像进行预处理
                                     1 / 255.0, # 图像各通道数值的缩放比例
                                     (416, 416), # 输出图像的空间尺寸
                                     swapRB=True,#交换RB通道,默认读取通道顺序为bgr通道
                                     crop=False)#图像裁剪
        self.net.setInput(blob)

        start = time.time()
        layerOutputs = self.net.forward(self.ln) # 输出层结果向前传播，得到检测结果：3个通道的shape中存放检测结果
        end = time.time()

        boxes = []  # 用于检测框坐标
        confidences = []  # 用于存放置信度值
        classIDs = []  # 用于识别的类别序号

        (H, W) = frame_in.shape[:2]#获取原图像尺寸
        # 逐层遍历网络获取输出
        for output in layerOutputs:
            # loop over each of the detections
            for detection in output:#对每个输出层中的每个检测框循环
                # extract the class ID and confidence (i.e., probability)
                # of the current object detection
                scores = detection[5:] # detection=[x,y,h,w,c,classes,……]中心点坐标，图像尺寸，置信度，类别概率
                classID = np.argmax(scores) # np.argmax反馈最大值的索引
                confidence = scores[classID]

                # 过滤低置信度值的检测结果
                if confidence > self.filter_confidence:
                    box = detection[0:4] * np.array([W, H, W, H])
                    (centerX, centerY, width, height) = box.astype("int")

                    #  边框的左上角
                    x = int(centerX - (width / 2))
                    y = int(centerY - (height / 2))

                    # 更新标记框、置信度值、类别列表
                    boxes.append([x, y, int(width), int(height)])
                    confidences.append(float(confidence))
                    classIDs.append(classID)

        # 使用NMS去除重复的标记框
        idxs = cv2.dnn.NMSBoxes(boxes, confidences, self.filter_confidence, self.threshold_prob)

        dets = []
        if len(idxs) > 0:
            # 遍历索引得到检测结果
            for i in idxs.flatten():# 降维
                (x, y) = (boxes[i][0], boxes[i][1])
                (w, h) = (boxes[i][2], boxes[i][3])
                dets.append([x, y, x + w, y + h, confidences[i], classIDs[i]])

        np.set_printoptions(formatter={'float': lambda x: "{0:0.3f}".format(x)})#控制输出结果的精度
        dets = np.asarray(dets)

        # 使用sort算法，开始进行追踪

        tracks = self.tracker.update(dets)
        boxes = []  # 存放追踪到的标记框
        indexIDs = []  # 存放追踪到的序号
        cls_IDs = []  # 存放追踪到的类别

        for track in tracks:
            boxes.append([track[0], track[1], track[2], track[3]])
            indexIDs.append(int(track[4]))
            cls_IDs.append(int(track[5]))


        return dets, boxes, indexIDs, cls_IDs




